import matplotlib.pyplot as plt
import numpy as np


# 定义三维函数
def f(x, y):
    return x ** 2 + y ** 2


# 定义函数梯度
def grad_f(x, y):
    return 2 * x, 2 * y


# 梯度下降函数
def gradient_descent(grad, start_x, start_y, learning_rate, num_iterations):
    x, y = start_x, start_y
    xs, ys, zs = [], [], []
    for i in range(num_iterations):
        dx, dy = grad(x, y)
        x, y = x - learning_rate * dx, y - learning_rate * dy
        z = f(x, y)
        xs.append(x)
        ys.append(y)
        zs.append(z)

    return xs, ys, zs


# 初始点
start_x, start_y = 3.0, 3.0
# 梯度下降
xs, ys, zs = gradient_descent(grad_f, start_x, start_y, 0.1, 50)

# 创建图像和轴
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 创建x,y 的数据点
x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
x, y = np.meshgrid(x, y)

z = f(x, y)

# 绘制表面
ax.plot_surface(x, y, z, cmap='viridis', alpha=0.7)

# 绘制梯度下降的点
ax.scatter(xs, ys, zs, color='r', s=50)

# 绘制点的连线
for i in range(len(xs) - 1):
    ax.plot([xs[i], xs[i + 1]], [ys[i], ys[i + 1]], [zs[i], zs[i + 1]], 'r-')

# 设置标签和标题
ax.set_xlabel('X  axis')
ax.set_ylabel('Y  axis')
ax.set_zlabel('z  axis')

ax.set_title('Gradient Descent on 3D  Surface')

# 显示图像
plt.show()
